"""
Some preprocessing and postprocessing
"""
import numpy as np
from numpy import ndarray
import cc3d


def keep_largest_component(input_mask: ndarray) -> ndarray:
    labels_out: ndarray = cc3d.connected_components(input_mask)
    # 索引值，索引出现的次数
    index_, count = np.unique(labels_out, return_counts=True)
    count_cp = count.copy()
    count_cp.sort()

    mark = index_[-1] + 1
    for c in count_cp[-2: -1]:
        # 取得相应次数的索引值
        ind = np.where(count == c)[0]
        print(f'index: {ind}, count: {c}')
        labels_out[labels_out == ind] = mark

    labels_out[labels_out != mark] = 0
    labels_out[labels_out == mark] = 1

    print(f'{input_mask.sum()}, {labels_out.sum()}')
    return labels_out.astype(dtype=input_mask.dtype)


def remove_small_object(input_mask: ndarray, limit: int = 100) -> ndarray:
    labels_out: ndarray = cc3d.connected_components(input_mask)
    # 索引值，索引出现的次数
    index_, count = np.unique(labels_out, return_counts=True)
    count_cp = count.copy()
    count_cp.sort()
    count_cp = filter(lambda x: x > limit, count_cp[:-1])  # 移除最大的连通域——背景

    mark = index_[-1] + 1
    for c in count_cp:
        # 取得相应次数的索引值
        # TODO 处理 count 相同的情况
        ind = np.where(count == c)[0]
        print(f'index: {ind}, count: {c}， {ind}， {mark}')
        labels_out[labels_out == ind] = mark

    labels_out[labels_out != mark] = 0
    labels_out[labels_out == mark] = 1

    print(f'{input_mask.sum()}, {labels_out.sum()}')
    return labels_out.astype(dtype=input_mask.dtype)
